Skip to content

Add Pruna AI library snippets and integration #1684

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Pruna AI library snippets and integration #1684

wants to merge 1 commit into from

Conversation

davidberenstein1957
Copy link

@davidberenstein1957 davidberenstein1957 commented Aug 12, 2025 •
edited
Loading

This update enhances support for Pruna AI, providing users with tailored code snippets for model integrations with Transformers and Diffusers.

  • Introduced a new library entry for Pruna AI in model-libraries.
  • Added main entry point and specific snippet generation functions for diffusers and transformers models.
  • Cleaned up whitespace inconsistencies in existing snippets.

TLDR: Pruna API normally mimics the Transformers and Diffusers API, so we can use PrunaModel.from_pretrained on top of pipelines or specific models. We re-use the underlying snippets for both the library and do some greedy replacements of certain part of the code snippets.

example

import torch
from diffusers import FluxFillPipeline
from diffusers.utils import load_image

image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")

pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
image = pipe(
    prompt="a white paper cup",
    image=image,
    mask_image=mask,
    height=1632,
    width=1232,
    guidance_scale=30,
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save(f"flux-fill-dev.png")

becomes

import torch
from pruna import PrunaModel
from diffusers.utils import load_image

image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")

pipe = PrunaModel.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
image = pipe(
    prompt="a white paper cup",
    image=image,
    mask_image=mask,
    height=1632,
    width=1232,
    guidance_scale=30,
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save(f"flux-fill-dev.png")

@davidberenstein1957
Add Pruna AI library snippets and integration
6287754 
- Introduced a new library entry for Pruna AI in model-libraries.
- Added main entry point and specific snippet generation functions for diffusers and transformers models.
- Cleaned up whitespace inconsistencies in existing snippets.

This update enhances support for Pruna AI, providing users with tailored code snippets for model integration.
@sharpenb
Copy link

This is super cool :) It looks good from the Pruna side. Loading models with this snippet would facilitate their compression to make them more efficient.

merveenoyan
merveenoyan reviewed Aug 14, 2025
View reviewed changes
Copy link
Contributor

@merveenoyan merveenoyan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks a lot, it's nice to see you here! interestingly this adds a lot of newlines, would be great to run linter

I think no need for a lot of comments to keep the code minimal :)

Vaibhavs10
Vaibhavs10 reviewed Aug 15, 2025
View reviewed changes
Copy link
Member

@Vaibhavs10 Vaibhavs10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR @davidberenstein1957 - let's wait for @Wauplin to give it a review, since this PR is a bit more involved!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Vaibhavs10 Vaibhavs10 Vaibhavs10 left review comments

@merveenoyan merveenoyan merveenoyan left review comments

@SBrandeis SBrandeis Awaiting requested review from SBrandeis SBrandeis is a code owner

@gary149 gary149 Awaiting requested review from gary149 gary149 is a code owner

@Wauplin Wauplin Awaiting requested review from Wauplin Wauplin is a code owner

@julien-c julien-c Awaiting requested review from julien-c julien-c is a code owner

@pcuenca pcuenca Awaiting requested review from pcuenca pcuenca is a code owner

@ngxson ngxson Awaiting requested review from ngxson ngxson is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@davidberenstein1957 @sharpenb @Vaibhavs10 @merveenoyan